import torch
import cn_clip.clip as clip
import torch.nn.functional as F
import torch.nn as nn


class LinearProbe(nn.Module):
    def __init__(self, config):
        super(LinearProbe, self).__init__()

        self.fc1 = nn.Linear(config["input_size"], config["hidden_size"])
        self.fc2 = nn.Linear(config["hidden_size"], config["hidden_size"])
        self.fc3 = nn.Linear(config["hidden_size"], config["num_classes"])

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


class Classifier(nn.Module):
    def __init__(self, config):
        super(Classifier, self).__init__()

        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'

        self.CLIP, self.img_processor = clip.load_from_name(
            "ViT-B-16", device='cuda', download_root=config["checkpoint_path"])
        self.tokenizer = clip.tokenize
        self.line_probe = LinearProbe(config)

    def forward(self, img, text):
        text = self.tokenizer(text).to(self.device)
        img = img.to(self.device)

        f_img = self.CLIP.encode_image(img)
        f_text = self.CLIP.encode_text(text)

        f_fuse = torch.cat((f_img, f_text), dim=1)
        predictions = self.line_probe(f_fuse)

        return predictions


